"""Fix unique constraints

Revision ID: 45724786402e
Revises: e22c96fb9959
Create Date: 2021-05-10 13:52:30.381832+00:00

"""
import itertools

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '45724786402e'
down_revision = 'e22c96fb9959'
branch_labels = None
depends_on = None


def ensure_unique_string(conn, table, column):
    values = set()
    for row in map(dict, conn.execute(f"SELECT * FROM {table}").fetchall()):
        if row[column] is not None:
            update = False
            if row[column] in values:
                update = True
                for i in itertools.count(1):
                    new_value = row[column] + str(i)
                    if new_value not in values:
                        row[column] = new_value
                        break

            if update:
                op.execute(f"UPDATE {table} SET {column} = ? WHERE id = ?", [row[column], row["id"]])

            values.add(row[column])


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    conn = op.get_bind()

    ensure_unique_string(conn, 'account_bsdgroups', 'bsdgrp_group')
    with op.batch_alter_table('account_bsdgroups', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_account_bsdgroups_bsdgrp_group'), ['bsdgrp_group'])

    ensure_unique_string(conn, 'account_bsdusers', 'bsdusr_username')
    with op.batch_alter_table('account_bsdusers', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_account_bsdusers_bsdusr_username'), ['bsdusr_username'])

    ensure_unique_string(conn, 'directoryservice_idmap_domain', 'idmap_domain_dns_domain_name')
    ensure_unique_string(conn, 'directoryservice_idmap_domain', 'idmap_domain_name')
    with op.batch_alter_table('directoryservice_idmap_domain', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_directoryservice_idmap_domain_idmap_domain_dns_domain_name'), ['idmap_domain_dns_domain_name'])
        batch_op.create_unique_constraint(batch_op.f('uq_directoryservice_idmap_domain_idmap_domain_name'), ['idmap_domain_name'])

    ensure_unique_string(conn, 'directoryservice_kerberoskeytab', 'keytab_name')
    with op.batch_alter_table('directoryservice_kerberoskeytab', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_directoryservice_kerberoskeytab_keytab_name'), ['keytab_name'])

    ensure_unique_string(conn, 'network_lagginterfacemembers', 'lagg_physnic')
    with op.batch_alter_table('network_lagginterfacemembers', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_network_lagginterfacemembers_lagg_physnic'), ['lagg_physnic'])

    ensure_unique_string(conn, 'services_iscsitarget', 'iscsi_target_alias')
    ensure_unique_string(conn, 'services_iscsitarget', 'iscsi_target_name')
    with op.batch_alter_table('services_iscsitarget', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_services_iscsitarget_iscsi_target_alias'), ['iscsi_target_alias'])
        batch_op.create_unique_constraint(batch_op.f('uq_services_iscsitarget_iscsi_target_name'), ['iscsi_target_name'])

    ensure_unique_string(conn, 'services_iscsitargetextent', 'iscsi_target_extent_naa')
    with op.batch_alter_table('services_iscsitargetextent', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_services_iscsitargetextent_iscsi_target_extent_naa'), ['iscsi_target_extent_naa'])

    ensure_unique_string(conn, 'storage_encrypteddisk', 'encrypted_provider')
    with op.batch_alter_table('storage_encrypteddisk', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_storage_encrypteddisk_encrypted_provider'), ['encrypted_provider'])

    ensure_unique_string(conn, 'storage_volume', 'vol_name')
    with op.batch_alter_table('storage_volume', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_storage_volume_vol_name'), ['vol_name'])

    ensure_unique_string(conn, 'system_acmednsauthenticator', 'name')
    with op.batch_alter_table('system_acmednsauthenticator', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_acmednsauthenticator_name'), ['name'])

    ensure_unique_string(conn, 'system_acmeregistration', 'directory')
    with op.batch_alter_table('system_acmeregistration', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_acmeregistration_directory'), ['directory'])

    ensure_unique_string(conn, 'system_certificate', 'cert_name')
    with op.batch_alter_table('system_certificate', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_certificate_cert_name'), ['cert_name'])

    ensure_unique_string(conn, 'system_certificateauthority', 'cert_name')
    with op.batch_alter_table('system_certificateauthority', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_certificateauthority_cert_name'), ['cert_name'])

    ensure_unique_string(conn, 'system_filesystem', 'identifier')
    with op.batch_alter_table('system_filesystem', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_filesystem_identifier'), ['identifier'])

    ensure_unique_string(conn, 'system_keyvalue', 'key')
    with op.batch_alter_table('system_keyvalue', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_keyvalue_key'), ['key'])

    ensure_unique_string(conn, 'system_migration', 'name')
    with op.batch_alter_table('system_migration', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_migration_name'), ['name'])

    ensure_unique_string(conn, 'system_tunable', 'tun_var')
    with op.batch_alter_table('system_tunable', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_system_tunable_tun_var'), ['tun_var'])

    ensure_unique_string(conn, 'truenas_enclosurelabel', 'encid')
    with op.batch_alter_table('truenas_enclosurelabel', schema=None) as batch_op:
        batch_op.create_unique_constraint(batch_op.f('uq_truenas_enclosurelabel_encid'), ['encid'])

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('truenas_enclosurelabel', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_truenas_enclosurelabel_encid'), type_='unique')

    with op.batch_alter_table('system_tunable', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_tunable_tun_var'), type_='unique')

    with op.batch_alter_table('system_migration', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_migration_name'), type_='unique')

    with op.batch_alter_table('system_keyvalue', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_keyvalue_key'), type_='unique')

    with op.batch_alter_table('system_filesystem', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_filesystem_identifier'), type_='unique')

    with op.batch_alter_table('system_certificateauthority', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_certificateauthority_cert_name'), type_='unique')

    with op.batch_alter_table('system_certificate', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_certificate_cert_name'), type_='unique')

    with op.batch_alter_table('system_acmeregistration', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_acmeregistration_directory'), type_='unique')

    with op.batch_alter_table('system_acmednsauthenticator', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_system_acmednsauthenticator_name'), type_='unique')

    with op.batch_alter_table('storage_volume', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_storage_volume_vol_name'), type_='unique')

    with op.batch_alter_table('storage_encrypteddisk', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_storage_encrypteddisk_encrypted_provider'), type_='unique')

    with op.batch_alter_table('services_iscsitargetextent', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_services_iscsitargetextent_iscsi_target_extent_naa'), type_='unique')

    with op.batch_alter_table('services_iscsitarget', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_services_iscsitarget_iscsi_target_name'), type_='unique')
        batch_op.drop_constraint(batch_op.f('uq_services_iscsitarget_iscsi_target_alias'), type_='unique')

    with op.batch_alter_table('network_lagginterfacemembers', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_network_lagginterfacemembers_lagg_physnic'), type_='unique')

    with op.batch_alter_table('directoryservice_kerberoskeytab', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_directoryservice_kerberoskeytab_keytab_name'), type_='unique')

    with op.batch_alter_table('directoryservice_idmap_domain', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_directoryservice_idmap_domain_idmap_domain_name'), type_='unique')
        batch_op.drop_constraint(batch_op.f('uq_directoryservice_idmap_domain_idmap_domain_dns_domain_name'), type_='unique')

    with op.batch_alter_table('account_bsdusers', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_account_bsdusers_bsdusr_username'), type_='unique')

    with op.batch_alter_table('account_bsdgroups', schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f('uq_account_bsdgroups_bsdgrp_group'), type_='unique')

    # ### end Alembic commands ###
